#include <mcp/server/resource_manager.h>
#include <mcp/json/value.h>
#include <mcp/error.h>
#include <mcp/types.h>
#include <iostream>
#include <sstream>
#include <algorithm>
#include <memory>
#include <regex>

namespace mcp {
namespace server {

// ResourceManager implementation
struct ResourceManager::Impl {
    std::unordered_map<std::string, ResourceRegistration> resources;
    std::vector<std::string> resource_order; // For maintaining insertion order

    // Thread safety
    mutable std::mutex mutex;
};

ResourceManager::ResourceManager() : impl_(std::make_unique<Impl>()) {}

ResourceManager::~ResourceManager() = default;

bool ResourceManager::register_resource(ResourceRegistration registration) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    if (impl_->resources.contains(registration.uri_template)) {
        return false; // Resource already exists
    }

    impl_->resources[registration.uri_template] = std::move(registration);
    impl_->resource_order.push_back(registration.uri_template);
    return true;
}

bool ResourceManager::unregister_resource(const std::string& uri_template) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->resources.find(uri_template);
    if (it == impl_->resources.end()) {
        return false; // Resource not found
    }

    impl_->resources.erase(it);
    impl_->resource_order.erase(
        std::remove(impl_->resource_order.begin(), impl_->resource_order.end(), uri_template),
        impl_->resource_order.end()
    );
    return true;
}

bool ResourceManager::has_resource(const std::string& uri_template) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    return impl_->resources.contains(uri_template);
}

std::vector<types::Resource> ResourceManager::list_resources() const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    std::vector<types::Resource> result;
    result.reserve(impl_->resource_order.size());

    for (const auto& uri_template : impl_->resource_order) {
        auto it = impl_->resources.find(uri_template);
        if (it != impl_->resources.end()) {
            const auto& registration = it->second;
            result.push_back(types::Resource{
                .uri = registration.uri_template,
                .name = registration.name,
                .description = registration.description,
                .mime_type = registration.mime_type
            });
        }
    }

    return result;
}

std::vector<types::ResourceTemplate> ResourceManager::list_resource_templates() const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    std::vector<types::ResourceTemplate> result;
    result.reserve(impl_->resource_order.size());

    for (const auto& uri_template : impl_->resource_order) {
        auto it = impl_->resources.find(uri_template);
        if (it != impl_->resources.end()) {
            const auto& registration = it->second;

            // Extract arguments from URI template
            std::vector<std::string> arguments = extract_template_arguments(uri_template);

            result.push_back(types::ResourceTemplate{
                .uri_template = registration.uri_template,
                .name = registration.name,
                .description = registration.description,
                .arguments = arguments
            });
        }
    }

    return result;
}

std::optional<types::Resource> ResourceManager::get_resource(const std::string& uri_template) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->resources.find(uri_template);
    if (it == impl_->resources.end()) {
        return std::nullopt;
    }

    const auto& registration = it->second;
    return types::Resource{
        .uri = registration.uri_template,
        .name = registration.name,
        .description = registration.description,
        .mime_type = registration.mime_type
    };
}

Result<types::ResourceContent> ResourceManager::read_resource(const std::string& uri,
                                                               Context<>& context) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    // Find matching template for the URI
    auto match = find_matching_template(uri);
    if (!match) {
        return Error::not_found(std::string("No matching resource template found for URI: ") + uri);
    }

    const auto& [registration, arguments] = *match;

    try {
        context.info(std::format("Reading resource: {} (template: {})", uri, registration.uri_template));

        // Call the resource handler
        auto result = registration.handler(uri, context);
        if (!result) {
            return result.error();
        }

        // Ensure the result has the correct URI
        result->uri = uri;

        // Validate the result
        if (result->contents.empty()) {
            context.warning(std::format("Resource '{}' returned empty content", uri));
        }

        return result;

    } catch (const std::exception& e) {
        return Error::internal_error(
            std::string("Exception in resource handler for '") + uri + "': " + e.what()
        );
    } catch (...) {
        return Error::internal_error(
            std::string("Unknown exception in resource handler for '") + uri + "'"
        );
    }
}

std::optional<std::pair<ResourceRegistration, std::unordered_map<std::string, std::string>>>
ResourceManager::find_matching_template(const std::string& uri) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    for (const auto& [uri_template, registration] : impl_->resources) {
        if (validate_uri(uri_template, uri)) {
            auto args = extract_arguments(uri_template, uri);
            if (args) {
                return std::make_pair(registration, std::move(*args));
            }
        }
    }

    return std::nullopt;
}

std::expected<std::unordered_map<std::string, std::string>, std::string>
ResourceManager::extract_arguments(const std::string& uri_template, const std::string& uri) {
    std::unordered_map<std::string, std::string> arguments;

    // Simple template matching - look for {arg} patterns
    std::regex arg_pattern(R"(\{([^}]+)\})");
    std::sregex_iterator template_it(uri_template.begin(), uri_template.end(), arg_pattern);
    std::sregex_iterator template_end;

    std::string pattern = uri_template;
    std::vector<std::string> arg_names;

    // Build regex pattern and extract argument names
    for (auto it = template_it; it != template_end; ++it) {
        std::string arg_name = (*it)[2].str();
        arg_names.push_back(arg_name);

        // Replace {arg} with (.+) in the pattern
        size_t pos = pattern.find((*it)[0].str());
        if (pos != std::string::npos) {
            pattern.replace(pos, (*it)[0].length(), "(.+)");
        }
    }

    // Convert to proper regex
    std::string regex_pattern = std::regex_replace(pattern, std::regex(R"([.^$*+?()\[\]{}|])"), R"(\$&)");
    std::regex uri_regex(regex_pattern);

    std::smatch matches;
    if (std::regex_match(uri, matches, uri_regex) && matches.size() == arg_names.size() + 1) {
        for (size_t i = 0; i < arg_names.size(); ++i) {
            arguments[arg_names[i]] = matches[i + 1].str();
        }
        return arguments;
    }

    return std::unexpected("URI does not match template pattern");
}

bool ResourceManager::validate_uri(const std::string& uri_template, const std::string& uri) {
    // Check if the basic structure matches
    size_t template_pos = 0;
    size_t uri_pos = 0;

    while (template_pos < uri_template.size() && uri_pos < uri.size()) {
        if (uri_template[template_pos] == '{') {
            // Start of argument - find closing brace
            size_t end_pos = uri_template.find('}', template_pos);
            if (end_pos == std::string::npos) {
                return false; // Unclosed brace
            }

            // Skip the argument in the URI
            template_pos = end_pos + 1;

            // Find the next literal character in the template
            while (template_pos < uri_template.size() && uri_template[template_pos] != '{' &&
                   uri_pos < uri.size() && uri[uri_pos] != uri_template[template_pos]) {
                uri_pos++;
            }
        } else {
            // Literal character match
            if (uri_template[template_pos] != uri[uri_pos]) {
                return false;
            }
            template_pos++;
            uri_pos++;
        }
    }

    // Both should be at the end
    return template_pos == uri_template.size() && uri_pos == uri.size();
}

void ResourceManager::clear() {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    impl_->resources.clear();
    impl_->resource_order.clear();
}

// Helper functions
std::vector<std::string> ResourceManager::extract_template_arguments(const std::string& uri_template) {
    std::vector<std::string> arguments;
    std::regex arg_pattern(R"(\{([^}]+)\})");
    std::sregex_iterator it(uri_template.begin(), uri_template.end(), arg_pattern);
    std::sregex_iterator end;

    for (; it != end; ++it) {
        arguments.push_back((*it)[2].str());
    }

    return arguments;
}

// Resource content helpers
types::ResourceContent create_text_resource_content(const std::string& text, const std::string& uri) {
    return types::ResourceContent{
        .contents = {std::make_shared<types::TextContent>(text)},
        .uri = uri
    };
}

types::ResourceContent create_image_resource_content(const std::string& image_data,
                                                      const std::string& mime_type,
                                                      const std::string& uri) {
    return types::ResourceContent{
        .contents = {std::make_shared<types::ImageContent>(image_data, mime_type)},
        .uri = uri
    };
}

types::ResourceContent create_mixed_resource_content(const std::vector<std::shared_ptr<types::ContentBlock>>& content_blocks,
                                                      const std::string& uri) {
    return types::ResourceContent{
        .contents = content_blocks,
        .uri = uri
    };
}

// URI validation helpers
bool validate_resource_uri(const std::string& uri) {
    if (uri.empty()) {
        return false;
    }

    // Basic URI validation - must have a scheme
    size_t colon_pos = uri.find(':');
    if (colon_pos == std::string::npos || colon_pos == 0) {
        return false;
    }

    std::string scheme = uri.substr(0, colon_pos);
    if (scheme.empty()) {
        return false;
    }

    // Check for common resource schemes
    static const std::vector<std::string> valid_schemes = {
        "file", "http", "https", "ftp", "data", "resource"
    };

    if (std::find(valid_schemes.begin(), valid_schemes.end(), scheme) == valid_schemes.end()) {
        // Allow custom schemes as long as they follow the rules
        for (char c : scheme) {
            if (!(std::isalnum(c) || c == '+' || c == '-' || c == '.')) {
                return false;
            }
        }
    }

    // Must have something after the scheme
    if (uri.size() <= colon_pos + 1) {
        return false;
    }

    return true;
}

std::string get_uri_scheme(const std::string& uri) {
    size_t colon_pos = uri.find(':');
    if (colon_pos == std::string::npos) {
        return "";
    }
    return uri.substr(0, colon_pos);
}

std::string get_uri_path(const std::string& uri) {
    size_t colon_pos = uri.find(':');
    if (colon_pos == std::string::npos || colon_pos + 1 >= uri.size()) {
        return "";
    }
    return uri.substr(colon_pos + 1);
}

// Resource template helpers
bool is_template_uri(const std::string& uri_template) {
    return uri_template.find('{') != std::string::npos &&
           uri_template.find('}') != std::string::npos;
}

std::string instantiate_template(const std::string& uri_template,
                                 const std::unordered_map<std::string, std::string>& arguments) {
    std::string result = uri_template;

    for (const auto& [key, value] : arguments) {
        std::string placeholder = "{" + key + "}";
        size_t pos = result.find(placeholder);
        while (pos != std::string::npos) {
            result.replace(pos, placeholder.length(), value);
            pos = result.find(placeholder, pos + value.length());
        }
    }

    return result;
}

// Resource execution context
class ResourceExecutionContext {
public:
    explicit ResourceExecutionContext(const std::string& uri, Context<>& context)
        : uri_(uri), context_(context) {
        start_time_ = std::chrono::steady_clock::now();
    }

    ~ResourceExecutionContext() {
        auto end_time = std::chrono::steady_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time_);
        context_.info(std::format("Resource '{}' read completed in {}ms", uri_, duration.count()));
    }

    void log_access() {
        context_.info(std::format("Accessing resource: {}", uri_));
    }

private:
    std::string uri_;
    Context<>& context_;
    std::chrono::steady_clock::time_point start_time_;
};

// Resource caching (simple implementation)
class ResourceCache {
public:
    struct CacheEntry {
        types::ResourceContent content;
        std::chrono::steady_clock::time_point timestamp;
        std::chrono::seconds ttl;
    };

    void put(const std::string& uri, const types::ResourceContent& content, std::chrono::seconds ttl) {
        std::lock_guard<std::mutex> lock(mutex_);
        cache_[uri] = CacheEntry{
            content,
            std::chrono::steady_clock::now(),
            ttl
        };
    }

    std::optional<types::ResourceContent> get(const std::string& uri) {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = cache_.find(uri);
        if (it == cache_.end()) {
            return std::nullopt;
        }

        auto now = std::chrono::steady_clock::now();
        if (now - it->second.timestamp > it->second.ttl) {
            cache_.erase(it);
            return std::nullopt;
        }

        return it->second.content;
    }

    void clear() {
        std::lock_guard<std::mutex> lock(mutex_);
        cache_.clear();
    }

    void cleanup_expired() {
        std::lock_guard<std::mutex> lock(mutex_);
        auto now = std::chrono::steady_clock::now();
        for (auto it = cache_.begin(); it != cache_.end();) {
            if (now - it->second.timestamp > it->second.ttl) {
                it = cache_.erase(it);
            } else {
                ++it;
            }
        }
    }

private:
    std::unordered_map<std::string, CacheEntry> cache_;
    mutable std::mutex mutex;
};

// Global resource cache instance
ResourceCache& get_resource_cache() {
    static ResourceCache instance;
    return instance;
}

} // namespace server
} // namespace mcp